from PIL import Image
import math


def sobel_edge_detection(image_path, output_path):
    """
    使用Sobel算子对输入图像进行边缘检测。

    参数:
    image_path (str): 输入图像的文件路径。
    output_path (str): 输出边缘检测后图像的文件路径。
    """
    # 打开原始图像并转换为灰度图（以灰度图为例）
    image = Image.open(image_path).convert('L')
    # 获取图像的宽度和高度
    width, height = image.size
    # 创建新的空白图像，尺寸与原图像相同，模式也相同，用于存储边缘检测结果
    new_image = Image.new('L', (width, height))
    # 定义Sobel算子（水平方向和垂直方向）
    sobel_x = [[-1, 0, 1],
               [-2, 0, 2],
               [-1, 0, 1]]
    sobel_y = [[-1, -2, -1],
               [0, 0, 0],
               [1, 2, 1]]
    # 遍历图像的每个像素点（避开边缘部分，边缘可后续单独处理或者忽略）
    for y in range(1, height - 1):
        for x in range(1, width - 1):
            # 计算水平方向梯度
            gradient_x = 0
            for ky in range(3):
                for kx in range(3):
                    pixel = image.getpixel((x + kx - 1, y + ky - 1))
                    gradient_x += pixel * sobel_x[ky][kx]
            # 计算垂直方向梯度
            gradient_y = 0
            for ky in range(3):
                for kx in range(3):
                    pixel = image.getpixel((x + kx - 1, y + ky - 1))
                    gradient_y += pixel * sobel_y[ky][kx]
            # 计算梯度幅值（采用平方根法）
            gradient_magnitude = math.sqrt(gradient_x ** 2 + gradient_y ** 2)
            # 将梯度幅值归一化到0-255范围（方便显示和保存为图像）
            normalized_gradient = int(gradient_magnitude / (math.sqrt(2) * 255) * 255)
            # 将归一化后的梯度幅值设置到新图像对应的坐标位置
            new_image.putpixel((x, y), normalized_gradient)
    # 保存边缘检测后的图像
    new_image.save(output_path)
    return new_image

image_path = "D:/cangku/computer-image-project-design/test.jpg"
output_path = "D:/cangku/computer-image-project-design/test1.jpg"
sobel_edge_detection(image_path, output_path)